'''
Created on Nov 12, 2013

@author: mkiyer
'''
from flask import Flask, jsonify, request, session, g, redirect, url_for, \
     abort, render_template, flash, make_response, json

# imports
import pymongo
import random
from cStringIO import StringIO
import matplotlib
matplotlib.use('Agg')
import re

# local imports
from ssea.base import Config, Metadata, SampleSet, Result
from ssea.countdata import BigCountMatrix
import detailedreport

# configuration
DATABASE = 'mongodb://localhost:27017/'
DEBUG = True
SECRET_KEY = 'development key'
#SECRET_KEY = '\x88\xa4EW\x8b\x98Hy\xbdwE\xd8?w%\x99LK\xa4+\xf89n\x91'
USERNAME = 'admin'
PASSWORD = 'default'

# limit number of documents returned by autocomplete queries 
SEARCH_AUTOCOMPLETE_LIMIT = 100

# create app
app = Flask(__name__)
app.config.from_object(__name__)

IGNORE_DATABASES = set(['local'])

class Study:
    ROW_META = 'metadata'
    COL_META = 'samples'
    CONFIG = 'config'
    SAMPLE_SETS = 'sample_sets'
    RESULTS = 'reports'    

    def __init__(self, db):
        self.db = db
        self.name = db.name
        self.config = Config.from_dict(self.db[Study.CONFIG].find_one())

def init_client():
    return pymongo.MongoClient(app.config['DATABASE'])

def init_studies(client):
    studies = {}
    for name in client.database_names():
        if name in IGNORE_DATABASES:
            continue
        db = client[name]
        if not Study.CONFIG in db.collection_names():
            #app.logger.debug('Skipping db %s' % (name))
            continue
        s = Study(db)
        studies[s.name] = s
    return studies

def get_client():
    if not hasattr(g, 'client'):
        g.client = init_client()
    return g.client

def get_studies():
    client = get_client()
    if not hasattr(g, 'studies'):
        g.studies = init_studies(client)
    return g.studies

@app.teardown_appcontext
def teardown(exception):
    if hasattr(g, 'studies'):
        # cleanup
        pass
    if hasattr(g, 'client'):
        g.client.close()

def get_mongo_regex(search_text):
    # simply search or OR search for strings with spaces
    return re.compile(re.escape(search_text), re.IGNORECASE)

@app.route('/heatmap')
def heatmap():
    return render_template('heatmap.html')

# @app.route('/search', methods=['POST'])
# def search():
#     t_id = request.form['t_id']
#     ss_id = request.form['ss_id']
#     app.logger.debug(t_id)
#     app.logger.debug(ss_id)
#     #flash('You entered something')
#     #return redirect(url_for('index'))
#     return query(t_id, ss_id)
#     #response = make_response(img_str)
#     #response.mimetype = 'image/png'
#     #return response
#     #img_str = query(t_id, ss_id)
#     #return img_str

# def connect_db(name='ssea'):
#     #return pymongo.MongoClient(app.config['DATABASE'])
#     client = pymongo.MongoClient(app.config['DATABASE'])
#     db = client.test
#     row_metadata = db[name + '.metadata']
#     col_metadata = db[name + '.samples']
#     sample_sets = db[name + '.sample_sets']
#     config = db[name + '.config']
#     reports = db[name + '.reports']
#     dbs = {'row_meta':row_metadata, 'col_meta':col_metadata, 'ss':sample_sets, 'config':config, 'reports':reports}
#     return dbs
# 
# def query(t_name, ss_name):  
#     colls = get_db()
#     ss = colls['ss'].find_one({'name': ss_name})
#     ss_id = ss['_id']
#     transcript = colls['row_meta'].find_one({'name': t_name})
#     transcript_id = transcript['_id']
#     # query for sample set and build SampleSet object
#     sample_set = SampleSet(**ss)
#     # query sample ids that are in sample set and build Metadata objects    
#     colmeta = list(Metadata.from_dict(x) for x in colls['col_meta'].find().sort('_id',1))   
#     # query t_id and build Metadata object
#     rowmeta = list(Metadata.from_dict(x) for x in colls['row_meta'].find().sort('_id',1))   
#     # read config json
#     runconfig = Config.from_dict(colls['config'].find_one())   
#     # query t_id and ss_id and get result
#     result = colls['reports'].find_one({'ss_id':ss_id, 't_id':transcript_id})   
#     result = Result.from_dict(result)
#     app.logger.debug(result)
#     # open matrix
#     bcm = BigCountMatrix.open(runconfig.matrix_dir)
#     # create detailed report
#     html_string = detailedreport.create_detailed_report(result, rowmeta[result.t_id], colmeta, sample_set, bcm, runconfig)
#     # cleanup
#     bcm.close()
#     return html_string
# 
# @app.route('/search', methods=['POST'])
# def search():
#     t_id = request.form['t_id']
#     ss_id = request.form['ss_id']
#     app.logger.debug(t_id)
#     app.logger.debug(ss_id)
#     #flash('You entered something')
#     #return redirect(url_for('index'))
#     return query(t_id, ss_id)
#     #response = make_response(img_str)
#     #response.mimetype = 'image/png'
#     #return response
#     #img_str = query(t_id, ss_id)
#     #return img_str

@app.route('/json_request_heatmap_data', methods=['POST'])
def json_request_heatmap_data():
    '''
    request is a json object with the following fields:
    
    new_rows: list of row ids (transcript ids) to be added to heatmap
    new_cols: list of column ids (sample set ids) to be added to heatmap
    rows: list of row ids currently in heatmap (not including new)
    cols: list of col ids currently in heatmap (not including new)
    value: field name of heatmap value (e.g. 'global_nes')
    
    the request needs to return a json object with the following data:
    
    rows: list of row metadata objects with _id field plus optional fields
    cols: list of col metadata object with _id field plus optional fields
    data: list of result objects each with 'row', 'col', and any number of value fields
    '''
    d = request.get_json()
    #app.logger.debug("heatmap data request: %s" % (repr(d)))
    study_name = d['study']
    new_row_ids = map(int, d['new_rows'])
    new_col_ids = map(int, d['new_cols'])
    all_row_ids = map(int, d['rows'])
    all_col_ids = map(int, d['cols'])
    value_field = str(d['value'])
    all_row_ids.extend(new_row_ids)
    all_col_ids.extend(new_col_ids)
    # lookup study name
    studies = get_studies()
    if study_name not in studies:
        # TODO: return error
        raise KeyError('study %s not in studies' % (study_name))
    db = studies[study_name].db
    # request data for new rows and columns
    # TODO: make this more intelligent
    # TODO: issue entire thing as one big batch? but make sure to only 
    # return data that the client needs
    coll = db[Study.RESULTS]
    data = {}
    for row_id in new_row_ids:
        for r in coll.find({'t_id': row_id, 'ss_id': {'$in': all_col_ids}},
                           {'t_id': 1, 'ss_id': 1, value_field: 1}):
            k = (r['t_id'], r['ss_id'])
            datum = {'row': r['t_id'], 
                     'col': r['ss_id'], 
                     'value': r[value_field]}
            data[k] = datum 
    for col_id in new_col_ids:
        for r in coll.find({'t_id': {'$in': all_row_ids}, 'ss_id': col_id},
                           {'t_id': 1, 'ss_id': 1, value_field: 1}):
            k = (r['t_id'], r['ss_id'])
            datum = {'row': r['t_id'], 
                     'col': r['ss_id'], 
                     'value': r[value_field]}
            data[k] = datum
    #app.logger.debug('returning rows: %s cols: %s data: %s' % (repr(new_rows), repr(new_cols), repr(data)))
    return jsonify(data=data.values())



@app.route('/_request_study_names', methods=['GET'])
def request_study_names():
    names = [s.name for s in get_studies()]
    return jsonify(names=names)  

@app.route('/json_search_sample_set', methods=['POST'])
def json_search_sample_set():
    '''
    request that sends json data with the fields:
    'study': name of study
    'q': input text in sample set search box
    
    returns json with the following fields:
    'id': sample set id
    'name': sample set name
    'desc': sample set description
    '''
    # parse the incoming json request
    d = request.get_json()
    study_name = d['study']
    query_string = d['q']
    # check for study name
    studies = get_studies()
    if study_name not in studies:
        # TODO: return error response
        raise KeyError('study %s not in studies' % (study_name))
    coll = studies[study_name].db[Study.SAMPLE_SETS]
    # search for query text (case insensitive)
    query = re.compile(re.escape(query_string), re.IGNORECASE)
    spec = {'name': query}
    response_fields = ['_id', 'name', 'desc']    
    results = []
    for d in coll.find(spec, response_fields, limit=SEARCH_AUTOCOMPLETE_LIMIT):
        result = dict((x,d[x]) for x in response_fields)
        results.append(result)
    return jsonify(results=results)

@app.route('/download_current_image', methods=['POST'])
def download_current_image():
    '''
    form submitted with two fields
    output_format: should be 'svg', 'png', or 'pdf'
    data: svg data to render to an image
    '''
    # form submitted with two fields
    # output_format: should
    output_format = request.form['output_format']
    data = request.form['data']
    if output_format not in ('svg', 'png', 'pdf'):
        # TODO: return error
        app.logger.error('Unrecognized image format %s' % (output_format))
    # svg output
    if output_format == 'svg':
        # no need to render svg simply return data to the user
        response = make_response(data)
        response.mimetype = 'image/svg+xml'
        response.headers['Content-Type'] = 'image/svg+xml'
        response.headers['Content-Disposition'] = 'attachment; filename=image.svg'
        return response
    else:
        app.logger.debug('not yet supported')

# ## PDF/PNG output
# elsif ($output_format eq "pdf" || $output_format eq "png") {
#     # Create temporary files (will be used with 'rsvg-convert')
#     my (undef, $input_file) = tempfile("d3export.svg.XXXXXXX", OPEN=>0, TMPDIR=>1, UNLINK=>1);
#     my (undef, $output_file) = tempfile("d3export.out.XXXXXXX", OPEN=>0, TMPDIR=>1, UNLINK=>1);
# 
#     # Write  the SVG data to a temporary file
#     write_file( $input_file, $data );
# 
#     my $zoom = ($output_format eq "png")?10:1;
# 
#     # Run "rsvg-convert", create the PNG/PDF file.
#     system("rsvg-convert -o '$output_file' -z '$zoom' -f '$output_format' '$input_file'");
# 
#     # Read the binary output (PDF/PNG) file.
#     my $pdf_data = read_file( $output_file, {binmode=>':raw'});
# 
#     ## All is fine, send the data back to the user
#     my $mime_type = ($output_format eq "pdf")?"application/x-pdf":"image/png";
#     print header(-type=>$mime_type,
#              -attachment=>"d3js_export_demo.$output_format");
#     print $pdf_data;
# 
#     exit(0);
# }
#     

@app.route('/studies', methods=['GET', 'POST'])
def request_studies():
    app.logger.debug('tried to fetch studies')
    studies = [{'_id': s} for s in get_studies()]
    return jsonify(results=studies)

@app.route('/samplesets', methods=['POST'])
def request_sample_sets():
    '''
    json request object with fields:
        'study': name of study
    
    returns json 'results' with list of objects each with attributes:
        '_id': sample set id
        'name': sample set name
        'desc': sample set description
    '''
    # parse the incoming json request
    d = request.get_json()
    study_name = d['study']
    app.logger.debug('sample sets for study %s' % (study_name))
    # check for study name
    studies = get_studies()
    if study_name not in studies:
        # TODO: return error response
        raise KeyError('study %s not in studies' % (study_name))
    coll = studies[study_name].db[Study.SAMPLE_SETS]
    # get all sample sets    
    fields = ['_id', 'name', 'desc']    
    results = []
    for d in coll.find(fields=fields):
        result = dict((x,d[x]) for x in fields)
        results.append(result)
    return jsonify(results=results)

@app.route('/transcripts', methods=['POST'])
def request_transcripts():
    '''
    json request object with fields:
        'study': name of study,
        'query': search string
    
    returns json 'results' with list of objects each with attributes:
        '_id': sample set id
        'name': sample set name
        'desc': sample set description
    '''
    # parse the incoming json request
    d = request.get_json()
    study_name = d['study']
    query_string = d['query']
    # check for study name
    studies = get_studies()
    if study_name not in studies:
        # TODO: return error response
        raise KeyError('study %s not in studies' % (study_name))
    coll = studies[study_name].db[Study.ROW_META]
    # search for query text (case insensitive)
    query = re.compile(re.escape(query_string), re.IGNORECASE)
    search_fields = ['name', 'category', 'locus', 'nearest_gene_names', 
                     'nearest_gene_ids']
    # Do a regex $or query and include each field
    spec = []
    for s in search_fields:
        spec.append({s: query})
    spec = {"$or": spec}
    results = []
    for d in coll.find(spec, limit=SEARCH_AUTOCOMPLETE_LIMIT):
        results.append(d)
    return jsonify(results=results)

@app.route('/index')
def index():
    return render_template('index.html', studies=get_studies())

@app.route('/')
def home():
    return render_template('home.html')


if __name__ == '__main__':
    app.run(debug=DEBUG)